import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import torch
from torchvision import transforms
import torch.nn.functional as func
from timm import create_model

num_classes = 4
model_name = 'tf_efficientnetv2_b2'
# classes = ['impressionism', 'realism', 'surrealism']  # 3
classes = ['expressionism', 'impressionism', 'realism', 'surrealism']  # 4
# classes = ['baroque', 'expressionism', 'impressionism', 'realism', 'surrealism']  # 5

# 加载模型
device = torch.device("mps")
model = create_model(model_name, pretrained=True)

model.load_state_dict(torch.load('Trained_Models/tf_efficientnetv2_b2_CL:4_VA:67.23_TA:94.91_23-12-05_14:02.pth',
                                 map_location=device))
model = model.to(device)
model.eval()

# 创建Tkinter窗口
window = tk.Tk()
window.title('Image Classifier')

# 获取屏幕的宽度和高度
screen_width = window.winfo_screenwidth()
screen_height = window.winfo_screenheight()

# 计算窗口应该位于的位置
x = (screen_width / 2) - (255 / 2)
y = (screen_height / 2) - (255 / 2)

# 设置窗口的位置，使窗口出现在屏幕中央
window.geometry("+%d+%d" % (x, y))

# 创建一个画布来显示图像
canvas = tk.Canvas(window, width=255, height=255)
canvas.pack()


# 创建一个按钮，当点击时，它会打开一个文件对话框来选择图像
def select_image():
    file_path = filedialog.askopenfilename()
    if not file_path:
        result_label.config(text='no input')
        return
    with open(file_path, 'rb') as f:
        image = Image.open(f).copy()

    max_size = (255, 255)
    image.thumbnail(max_size, Image.BICUBIC)
    tk_image = ImageTk.PhotoImage(image)

    image_width, image_height = image.size
    canvas_width = canvas.winfo_width()
    canvas_height = canvas.winfo_height()
    x = (canvas_width - image_width) / 2
    y = (canvas_height - image_height) / 2

    canvas.create_image(x, y, anchor='nw', image=tk_image)
    canvas.image = tk_image

    if not isinstance(image, Image.Image):
        transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Lambda(lambda x: x.convert('RGB')),
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x.float())
        ])
    else:
        transform = transforms.Compose([
            transforms.Lambda(lambda x: x.convert('RGB')),
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x.float())
        ])
    image = transform(image).unsqueeze(0).to(device)

    output = model(image)
    _, predicted = torch.max(output.data, 1)

    probabilities = func.softmax(output, dim=1)
    max_probability = torch.max(probabilities).item()

    result_label.config(
        text=f'Predicted class: {classes[predicted.item()]}, Probability: {max_probability * 100:.2f}%')


select_button = tk.Button(window, text='Select Image', command=select_image)
select_button.pack()

result_label = tk.Label(window, text='')
result_label.pack()

# 运行Tkinter事件循环
window.mainloop()
